import sys
import platform
import subprocess
import re


def check_python_version():
    """
    检查Python版本
    """
    version = sys.version_info
    print(f"Python版本: {version.major}.{version.minor}.{version.micro}")
    # PyTorch至少需要Python 3.6或更高版本
    if version.major < 3 or version.minor < 6:
        print("警告: 您的Python版本低于3.6，可能无法安装最新版本的PyTorch")
        return False
    return True


def check_operating_system():
    """
    检查操作系统信息
    """
    os_name = platform.system()
    os_version = platform.version()
    print(f"操作系统: {os_name}")
    print(f"系统版本: {os_version}")
    
    # PyTorch支持的主要操作系统
    supported_os = ['Windows', 'Linux', 'Darwin']  # Darwin是macOS
    if os_name not in supported_os:
        print(f"警告: 您的操作系统 {os_name} 可能不受PyTorch官方支持")
        return False
    return True


def check_gpu_windows():
    """
    在Windows系统上检查GPU信息
    """
    try:
        # 使用wmic命令检查GPU信息
        result = subprocess.run(['wmic', 'path', 'win32_VideoController', 'get', 'Name'],
                                capture_output=True, text=True, check=True)
        gpu_info = result.stdout.strip()
        
        if 'NVIDIA' in gpu_info:
            print("检测到NVIDIA GPU，支持安装GPU版本的PyTorch")
            # 尝试获取NVIDIA驱动版本
            try:
                nvcc_result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
                if nvcc_result.returncode == 0:
                    # 提取CUDA版本信息
                    cuda_version_match = re.search(r'CUDA Version: ([0-9.]+)', nvcc_result.stdout)
                    if cuda_version_match:
                        cuda_version = cuda_version_match.group(1)
                        print(f"CUDA版本: {cuda_version}")
                        return True, True, cuda_version
                    else:
                        print("无法确定CUDA版本，但检测到NVIDIA GPU")
                        return True, True, None
                else:
                    print("未找到nvidia-smi命令，可能需要安装NVIDIA驱动")
                    return True, False, None
            except FileNotFoundError:
                print("未找到nvidia-smi命令，可能需要安装NVIDIA驱动")
                return True, False, None
        else:
            print(f"检测到GPU，但不是NVIDIA GPU: {gpu_info}")
            print("只能安装CPU版本的PyTorch")
            return True, False, None
    except Exception as e:
        print(f"检查GPU信息时出错: {e}")
        print("默认假设只支持CPU版本的PyTorch")
        return True, False, None


def check_gpu_linux():
    """
    在Linux系统上检查GPU信息
    """
    try:
        # 检查是否有NVIDIA GPU
        result = subprocess.run(['lspci', '|', 'grep', '-i', 'vga'],
                                shell=True, capture_output=True, text=True)
        gpu_info = result.stdout.strip()
        
        if 'NVIDIA' in gpu_info:
            print("检测到NVIDIA GPU，支持安装GPU版本的PyTorch")
            # 尝试获取NVIDIA驱动版本
            try:
                nvcc_result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
                if nvcc_result.returncode == 0:
                    # 提取CUDA版本信息
                    cuda_version_match = re.search(r'CUDA Version: ([0-9.]+)', nvcc_result.stdout)
                    if cuda_version_match:
                        cuda_version = cuda_version_match.group(1)
                        print(f"CUDA版本: {cuda_version}")
                        return True, True, cuda_version
                    else:
                        print("无法确定CUDA版本，但检测到NVIDIA GPU")
                        return True, True, None
                else:
                    print("未找到nvidia-smi命令，可能需要安装NVIDIA驱动")
                    return True, False, None
            except FileNotFoundError:
                print("未找到nvidia-smi命令，可能需要安装NVIDIA驱动")
                return True, False, None
        else:
            print(f"检测到GPU，但不是NVIDIA GPU或无法识别: {gpu_info}")
            print("只能安装CPU版本的PyTorch")
            return True, False, None
    except Exception as e:
        print(f"检查GPU信息时出错: {e}")
        print("默认假设只支持CPU版本的PyTorch")
        return True, False, None


def check_gpu_macos():
    """
    在macOS系统上检查GPU信息
    """
    try:
        # 使用system_profiler命令检查GPU信息
        result = subprocess.run(['system_profiler', 'SPDisplaysDataType'],
                                capture_output=True, text=True, check=True)
        gpu_info = result.stdout
        
        if 'NVIDIA' in gpu_info:
            print("检测到NVIDIA GPU")
            # macOS上的CUDA支持有限
            print("注意: 最新版本的PyTorch在macOS上主要支持Metal加速，而非CUDA")
            return True, False, None  # 尽管有NVIDIA GPU，但macOS上建议使用CPU版本或Metal支持的版本
        elif 'AMD' in gpu_info or 'Intel Iris' in gpu_info or 'Apple M' in gpu_info:
            print("检测到非NVIDIA GPU")
            print("建议安装CPU版本的PyTorch或支持Metal加速的版本")
            return True, False, None
        else:
            print("未检测到独立GPU信息")
            print("只能安装CPU版本的PyTorch")
            return True, False, None
    except Exception as e:
        print(f"检查GPU信息时出错: {e}")
        print("默认假设只支持CPU版本的PyTorch")
        return True, False, None


def check_gpu():
    """
    根据操作系统检查GPU信息
    """
    os_name = platform.system()
    
    if os_name == 'Windows':
        return check_gpu_windows()
    elif os_name == 'Linux':
        return check_gpu_linux()
    elif os_name == 'Darwin':  # macOS
        return check_gpu_macos()
    else:
        print(f"不支持的操作系统: {os_name}")
        return False, False, None


def check_pytorch_compatibility():
    """
    综合检查当前系统是否支持PyTorch安装
    """
    print("========== PyTorch兼容性检查 ==========")
    
    # 检查Python版本
    python_ok = check_python_version()
    
    # 检查操作系统
    os_ok = check_operating_system()
    
    # 检查GPU
    gpu_supported, cuda_available, cuda_version = check_gpu()
    
    print("\n========== 兼容性检查结果 ==========")
    
    if python_ok and os_ok:
        print("✓ 当前系统基本支持安装PyTorch")
        
        # 给出安装建议
        print("\n安装建议:")
        if cuda_available and cuda_version:
            # 提取主要CUDA版本号（例如11.8 -> 11）
            major_cuda_version = cuda_version.split('.')[0]
            print(f"- 您的系统支持CUDA {cuda_version}，建议安装对应版本的GPU版PyTorch")
            print(f"  安装命令示例: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu{major_cuda_version}")
        elif gpu_supported and not cuda_available:
            print("- 检测到GPU，但CUDA不可用或未正确安装")
            print("  1. 您可以选择安装CPU版本的PyTorch")
            print("  2. 或者先安装NVIDIA驱动和CUDA，然后安装GPU版本的PyTorch")
            print("  CPU版本安装命令: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu")
        else:
            print("- 建议安装CPU版本的PyTorch")
            print("  安装命令: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu")
            
        print("\n注意：为了获得最准确的安装命令，请访问PyTorch官方网站：https://pytorch.org/")
        return True
    else:
        print("✗ 当前系统可能不支持安装PyTorch或需要升级")
        if not python_ok:
            print("- 请升级Python到3.6或更高版本")
        if not os_ok:
            print("- 您的操作系统可能不受PyTorch官方支持")
        return False


if __name__ == "__main__":
    check_pytorch_compatibility()